import { StoreApi, UseBoundStore } from "zustand";

interface State<T extends object> {
  (
    getData?: "" | keyof T,
    ...rest: Array<keyof T & string>
  ): [
    T,
    (
      partial: T | Partial<T> | ((state: T) => T | Partial<T>),
      replace?: boolean | undefined
    ) => void,
  ];
  getState: () => T;
  setState: (
    partial: T | Partial<T> | ((state: T) => T | Partial<T>),
    replace?: boolean | undefined
  ) => void;
}

export const createSelector = <S extends object>(
  _store: UseBoundStore<StoreApi<S>>
): State<S> => {
  const func: State<S> = (getData, ...rest) => {
    if (getData === undefined) {
      return [_store(), _store.setState];
    }

    const obj = { ..._store.getState() };
    if (getData === "") return [obj, _store.setState];

    return [
      _store((state) =>
        [getData, ...rest].reduce(
          (acc, key) => {
            acc[key] = state[key];
            return acc;
          },
          { ...obj }
        )
      ),
      _store.setState,
    ];
  };

  func.getState = _store.getState;
  func.setState = _store.setState;

  return func;
};

interface StateSelector<T extends object> {
  (getData?: keyof T, ...rest: Array<keyof T & string>): T;
  getState: () => T;
  setState: (
    partial: T | Partial<T> | ((state: T) => T | Partial<T>),
    replace?: boolean | undefined
  ) => void;
}
export const stateSelector = <T extends object>(
  _store: UseBoundStore<StoreApi<T>>
): StateSelector<T> => {
  const func: StateSelector<T> = (getData, ...rest) => {
    if (getData === undefined) {
      return _store();
    }

    const obj = { ..._store.getState() };
    return _store((state) =>
      [getData, ...rest].reduce(
        (acc, key) => {
          acc[key] = state[key];
          return acc;
        },
        { ...obj }
      )
    );
  };

  func.getState = _store.getState;
  func.setState = _store.setState;

  return func;
};
